'''
整理json文件
'''

import os
import numpy as np

from tool import read_json, write_json, json_name_deal, response_deal


def main(save_path:str="ckpt", output_path:str="tmp/result_v3.json", json_name_format:list=["country-mode", "model", "exp-index"]):
    # 读取ckpt目录下的所有文件
    files = os.listdir(save_path)

    all_data = {}
    # 遍历所有文件
    for file in files:
        # 读取文件
        data = read_json(f"{save_path}/{file}")

        if "en" in file:
            json_name_format = ["country-mode", "country-en", "model", "language", "exp-index"]
        else:
            json_name_format = ["country-mode", "country-zh", "model", "exp-index"]
        key = json_name_deal(file, json_name_format)
        # 整理数据
        if key["model"] not in all_data.keys():
            all_data[key["model"]] = np.zeros((50, 18)).tolist()

        response_data = response_deal(data["response"])
        for i in range(1, 19):
            all_data[key["model"]][int(key["exp-index"])-1][i-1] = response_data[f"{i}"]

    # 保存数据
    write_json(output_path, all_data)

if __name__ == "__main__":
    main(save_path="newckpt", output_path="tmp/result_newckpt.json",
         json_name_format=["country-mode", "country", "model", "exp-index"])
        